This PR addresses a memory accumulation issue observed during training with the ResGatedDynamicGNI model when persistent_workers=True is enabled in the DataLoader.

⚙️ Existing Problem

• The ResGatedDynamicGNI model performs per-forward random feature initialization for both node and edge features (new_x, new_edge_attr) on the GPU.

• When combined with persistent DataLoader workers, these per-batch random allocations are not released properly because:

  • Worker processes remain alive across epochs.
  • CUDA’s caching allocator retains fragmented memory blocks.

While setting persistent_workers=True can improve performance when input features remain constant throughout training (as noted in the Lightning documentation), it becomes problematic when input features are dynamically initialized in each forward pass. In such cases, the DataLoader workers retain these transient tensors in memory, expecting reuse across epochs. Since they are never reused, this leads to progressive GPU memory accumulation and can eventually cause out-of-memory (OOM) errors. See the related issue and logs here
.

Refer: https://lightning.ai/docs/pytorch/stable/advanced/speed.html#persistent-workers

🧠 Root Cause

persistent_workers=True keeps worker subprocesses alive between epochs, retaining CUDA contexts and cached memory allocations that the ResGatedDynamicGNI model reinitializes each forward pass.

🔧 Fix Implemented

• Enabled to set persistent_workers=False in all DataLoader through CLI for the ResGatedDynamicGNI model training.
  This ensures that:

  • Workers are restarted cleanly each epoch.
  • GPU and CPU memory are fully released after each epoch.
  • Memory fragmentation and accumulation are avoided.

• Default is set to True as before, ensuring no disruption for other existing pipelines